/*******************************************************************************
 * MPI Point-to-Point Communication Bandwidth Benchmark
 * 
 * This program measures the bidirectional bandwidth between two MPI ranks using
 * pipelined asynchronous send/receive operations. It specifically tests NUMA-aware
 * memory allocation to ensure optimal performance across different sockets.
 * 
 * Key features:
 * - Uses multiple buffered, pipelined communication for maximum throughput
 * - NUMA-aware memory allocation (rank 0 on socket 0, rank 1 on socket 1)
 * - Measures bidirectional bandwidth using MPI_Isend/MPI_Irecv
 * - Includes warmup phase to ensure stable measurements
 * 
 * Usage: mpirun -n 2 ./test_mpi_send_recv
 * 
 * Buffer size: 256MB
 * Number of iterations: 16
 * Pipeline depth: 16 buffers
 ******************************************************************************/

#include <stdio.h>
#include <stdlib.h>
#include <mpi.h>
#include <numa.h>
#include <string.h>
#include <sched.h>

#define BUFFER_SIZE (256*1024*1024) // 256MB buffer size
#define NUM_ITERATIONS 16
#define NUM_BUFFERS 16  // Increased pipeline depth

int main(int argc, char *argv[]) {
    int rank, size;
    double start_time, end_time, bandwidth;
    MPI_Request send_requests[NUM_BUFFERS];
    MPI_Request recv_requests[NUM_BUFFERS];
    void *buffers[NUM_BUFFERS];
    
    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);

    if (size != 2) {
        if (rank == 0) printf("This test requires exactly 2 MPI ranks\n");
        MPI_Finalize();
        return 1;
    }

    // Initialize NUMA
    if (numa_available() < 0) {
        printf("NUMA not available\n");
        MPI_Finalize();
        return 1;
    }

    // Get and display NUMA information
    int num_nodes = numa_num_configured_nodes();
    int cpu = sched_getcpu();
    int node = numa_node_of_cpu(cpu);
    
    printf("Rank %d: Running on CPU %d, NUMA node %d (Socket %d)\n", 
           rank, cpu, node, node/2);  // Assuming 2 NUMA nodes per socket
    MPI_Barrier(MPI_COMM_WORLD);  // Ensure clean output

    // Allocate multiple buffers on appropriate NUMA node based on rank
    // Rank 0 uses NUMA node 1 (Socket 0)
    // Rank 1 uses NUMA node 2 (Socket 1)
    for (int i = 0; i < NUM_BUFFERS; i++) {
        buffers[i] = numa_alloc_onnode(BUFFER_SIZE, rank == 0 ? 0 : 2);
        if (buffers[i] == NULL) {
            printf("Rank %d: Failed to allocate memory on NUMA node\n", rank);
            for (int j = 0; j < i; j++) {
                numa_free(buffers[j], BUFFER_SIZE);
            }
            MPI_Finalize();
            return 1;
        }
        // Touch each page to ensure proper NUMA allocation
        char *buf = (char *)buffers[i];
        for (size_t j = 0; j < BUFFER_SIZE; j += 4096) {
            buf[j] = 'a' + rank;
        }
    }

    MPI_Barrier(MPI_COMM_WORLD);

    // Warmup with async operations
    for (int i = 0; i < NUM_BUFFERS; i++) {
        if (rank == 0) {
            MPI_Isend(buffers[i], BUFFER_SIZE, MPI_CHAR, 1, i, MPI_COMM_WORLD, &send_requests[i]);
            MPI_Irecv(buffers[i], BUFFER_SIZE, MPI_CHAR, 1, i, MPI_COMM_WORLD, &recv_requests[i]);
        } else {
            MPI_Irecv(buffers[i], BUFFER_SIZE, MPI_CHAR, 0, i, MPI_COMM_WORLD, &recv_requests[i]);
            MPI_Isend(buffers[i], BUFFER_SIZE, MPI_CHAR, 0, i, MPI_COMM_WORLD, &send_requests[i]);
        }
    }
    MPI_Waitall(NUM_BUFFERS, send_requests, MPI_STATUSES_IGNORE);
    MPI_Waitall(NUM_BUFFERS, recv_requests, MPI_STATUSES_IGNORE);

    MPI_Barrier(MPI_COMM_WORLD);
    start_time = MPI_Wtime();

    // Main benchmark loop with pipelined async communication
    for (int iter = 0; iter < NUM_BUFFERS; iter++) {
        int buf_idx = iter % NUM_BUFFERS;

        // Post new async operations
        if (rank == 0) {
            MPI_Isend(buffers[buf_idx], BUFFER_SIZE, MPI_CHAR, 1, iter, MPI_COMM_WORLD, &send_requests[buf_idx]);
            MPI_Irecv(buffers[buf_idx], BUFFER_SIZE, MPI_CHAR, 1, iter, MPI_COMM_WORLD, &recv_requests[buf_idx]);
        } else {
            MPI_Irecv(buffers[buf_idx], BUFFER_SIZE, MPI_CHAR, 0, iter, MPI_COMM_WORLD, &recv_requests[buf_idx]);
            MPI_Isend(buffers[buf_idx], BUFFER_SIZE, MPI_CHAR, 0, iter, MPI_COMM_WORLD, &send_requests[buf_idx]);
        }
    }

    // Wait for all remaining operations to complete
    for (int i = 0; i < NUM_BUFFERS; i++) {
        MPI_Wait(&send_requests[i], MPI_STATUS_IGNORE);
        MPI_Wait(&recv_requests[i], MPI_STATUS_IGNORE);
    }

    end_time = MPI_Wtime();

    if (rank == 0) {
        double time_taken = end_time - start_time;
        // Calculate bandwidth in GB/s (accounting for bidirectional transfer)
        bandwidth = (2.0 * BUFFER_SIZE * NUM_BUFFERS) / (time_taken * 1024 * 1024 * 1024);
        printf("Time: %.3f seconds\n", time_taken);
        printf("Bandwidth: %.2f GB/s\n", bandwidth);
    }

    for (int i = 0; i < NUM_BUFFERS; i++) {
        numa_free(buffers[i], BUFFER_SIZE);
    }
    
    MPI_Finalize();
    return 0;
}

